// SPDX-FileCopyrightText: 2023 CASTest Corporation Limited
// SPDX-License-Identifier: LGPL-v3

#ifndef ICTEST_SPLITDECISIONTREE_H
#define ICTEST_SPLITDECISIONTREE_H

#include "atpg/ATPGDefine.h"
#include "atpg/LogicDrivenTable.h"
#include "netlist/PrimNetlist.h"

namespace ictest {

class SplitDecisionTree;

enum DecType { TEST_ROOT, PROPAGATION, GOOD_JUSTIFY, BAD_JUSTIFY };

// 决策树
class SplitSiblingTreeNode {
public:
  SplitSiblingTreeNode() {
    done_dec_ = false;
    initial_ = false;
    parent_ = nullptr;
    child_ = nullptr;
    sibling_ = nullptr;
    t_dec_ = TEST_ROOT;
    last_dec_ = 0;
  }

  virtual ~SplitSiblingTreeNode() = default;
  virtual bool another_dec() { return last_dec_ < max_dec_; }
  virtual int dec_type() { return t_dec_; }
  virtual bool initialized() { return initial_; }
  virtual void done_dec() { done_dec_ = true; }
  virtual void undone_dec() { done_dec_ = false; }
  virtual bool done() { return done_dec_; }
  virtual void init() {
    initial_ = false;
    max_dec_ = 0;
    last_dec_ = 0;
    imply_good_.clear();
    imply_bad_.clear();
  }
  virtual void set_imply(const std::vector<Gate *> &imply_good,
                         const std::vector<Gate *> &imply_bad) {
    imply_good_ = imply_good;
    imply_bad_ = imply_bad;
  }
  virtual void get_imply(std::vector<Gate *> &good, std::vector<Gate *> &bad) {
    good = imply_good_;
    bad = imply_bad_;
  }

public:
  bool initial_;
  bool done_dec_;
  int t_dec_;
  int last_dec_;
  int max_dec_;
  Gate *gate_;
  int val_;
  SplitSiblingTreeNode *parent_;
  SplitSiblingTreeNode *child_;
  SplitSiblingTreeNode *sibling_;
  std::vector<Gate *> imply_good_;
  std::vector<Gate *> imply_bad_;
};

class SplitJustifyNode : public SplitSiblingTreeNode {
public:
  SplitJustifyNode(SplitSiblingTreeNode *parent, Gate *gate, int val,
                   bool good) {
    parent_ = parent;
    gate_ = gate;
    val_ = val;
    if (good) {
      t_dec_ = GOOD_JUSTIFY;
    } else {
      t_dec_ = BAD_JUSTIFY;
    }
    max_dec_ = 0;
  }

  /**
   * @brief
   * init justify node
   */
  void init() override {
    SplitSiblingTreeNode::init();
    justify_.clear();
  }

  /**
   * @brief
   * get best jusitfy cube
   * @return best jusitfy cube
   * @attention
   * return by order adding to cubes
   */
  std::vector<std::pair<Gate *, int>> best_cube() {
    assert(last_dec_ < max_dec_);
    return justify_[last_dec_++];
  }

public:
  void calc_justify(LogicDrivenTable *driven_table,
                    std::map<Gate *, int> &gate2scanchainid,
                    std::vector<int> &good, std::vector<int> &bad,
                    std::vector<int> &fault_cone, int atpg_invokes,
                    std::vector<int> &is_falling_edge_dff,
                    std::vector<bool> &no_assign);

public:
  std::vector<std::vector<std::pair<Gate *, int>>> justify_;
};

class SplitPropagationNode : public SplitSiblingTreeNode {
public:
  SplitPropagationNode(SplitSiblingTreeNode *parent) {
    parent_ = parent;
    t_dec_ = PROPAGATION;
    max_dec_ = 0;
  }

  /**
   * @brief
   * init propagation node
   */
  void init() override {
    SplitSiblingTreeNode::init();
    dec_paths_.clear();
    path_justify_.clear();
  }

  /**
   * @brief
   * set decision path of current node
   * @param paths decision path
   */
  void set_dec_paths(const std::vector<std::pair<Gate *, Gate *>> &paths) {
    dec_paths_ = paths;
  }
/**
 * @brief
 * get best propagation cube and chosen path
 * @param chosen_path
 * @return best propagation cube
 * @attention
 * return by order adding to cubes
 */
  std::vector<std::pair<Gate *, int>>
  best_path(std::pair<Gate *, Gate *> &chosen_path) {
    assert(last_dec_ < max_dec_);
    int path = 0;
    int acc = 0;
    for (int p = 0; p < dec_paths_.size(); p++) {
      if ((last_dec_ - acc) >= path_justify_[p].size()) {
        acc += path_justify_[p].size();
      } else {
        path = p;
        break;
      }
    }

    assert(acc <= last_dec_);
    chosen_path = dec_paths_[path];
    auto res = path_justify_[path][last_dec_ - acc];
    last_dec_++;
    return res;
  }

public:
  void calc_propagation(LogicDrivenTable *driven_table,
                        std::map<Gate *, int> &gate2scanchainid,
                        std::vector<int> &good, std::vector<int> &bad,
                        std::vector<int> &fault_cone, int atpg_invokes,
                        std::vector<int> &is_falling_edge_dff,
                        std::vector<bool> &no_assign);

public:
  std::vector<std::pair<Gate *, Gate *>> dec_paths_;
  std::vector<std::vector<std::vector<std::pair<Gate *, int>>>> path_justify_;
};

class SplitDecisionTree {
public:
  SplitDecisionTree() {
    tg_ = new SplitSiblingTreeNode();
    tg_->child_ = nullptr;
    tg_->done_dec_ = true;
    tg_->max_dec_ = 0;
  }
  ~SplitDecisionTree() { delete_silb_child(tg_); }

  /**
   * @brief
   * get root node
   * @return root node ptr
   */
  SplitSiblingTreeNode *get_root() { return tg_; }
  // insert new decision nodes as child of root.
  static void insert_dec(SplitSiblingTreeNode *root,
                  const std::vector<SplitSiblingTreeNode *> &children);
  static void delete_silb_child(SplitSiblingTreeNode *root);
  SplitSiblingTreeNode *next_dec();

private:
  SplitSiblingTreeNode *tg_;
};

} // namespace ictest

#endif // ICTEST_SPLITDECISIONTREE_H
